package com.gitee.core.zipkin;

import com.gitee.core.dubbo.ZipkinConfig;
import com.gitee.core.wrapper.BodyReaderHttpServletRequestWrapper;
import com.github.kristofa.brave.*;
import com.github.kristofa.brave.http.*;
import org.springframework.web.servlet.handler.HandlerInterceptorAdapter;
import zipkin.reporter.AsyncReporter;
import zipkin.reporter.Reporter;
import zipkin.reporter.Sender;
import zipkin.reporter.okhttp3.OkHttpSender;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.net.URI;
import java.net.URISyntaxException;

public class ZipkinHandlerInterceptor extends HandlerInterceptorAdapter {


    static final String HTTP_SERVER_SPAN_ATTRIBUTE = ZipkinHandlerInterceptor.class.getName() + ".server-span";

    private final ServerRequestInterceptor requestInterceptor;
    private final ServerResponseInterceptor responseInterceptor;
    private final ServerSpanThreadBinder serverThreadBinder;
    private final SpanNameProvider spanNameProvider;

    ZipkinHandlerInterceptor() {
        String sendUrl = ZipkinConfig.getProperty("zipkin.address");
        String application = ZipkinConfig.getProperty("spring.application.name");
        Sender sender = OkHttpSender.create(sendUrl);
        Reporter<zipkin.Span> reporter = AsyncReporter.builder(sender).build();
        Brave brave = new Brave.Builder(application).reporter(reporter).build();

        this.requestInterceptor = brave.serverRequestInterceptor();
        this.responseInterceptor = brave.serverResponseInterceptor();
        this.serverThreadBinder = brave.serverSpanThreadBinder();
        this.spanNameProvider = new DefaultSpanNameProvider();
    }



    @Override
    public boolean preHandle(final HttpServletRequest request, final HttpServletResponse response, final Object handler) {
        String service = request.getRequestURI();
        service = service == null || service == "" ? "unkown" : service;
        BodyReaderHttpServletRequestWrapper requestWrapper = new BodyReaderHttpServletRequestWrapper(request);
        String data = requestWrapper.getBody() == null ? "{}" : requestWrapper.getBody();

        requestInterceptor.handle(new ZipkinRequestAdapter(new HttpServerRequest() {

            public String getHttpHeaderValue(String headerName) {
                return request.getHeader(headerName);
            }


            public URI getUri() {
                try {
                    return new URI(request.getRequestURI());
                } catch (URISyntaxException e) {
                    throw new RuntimeException(e);
                }
            }


            public String getHttpMethod() {
                String service = request.getRequestURI();
                service = service == null || service == "" ? "unkown" : service;
                return service;
            }
        }, service, data, spanNameProvider));

        return true;
    }

    @Override
    public void afterConcurrentHandlingStarted(final HttpServletRequest request, final HttpServletResponse response, final Object handler) {
        request.setAttribute(HTTP_SERVER_SPAN_ATTRIBUTE, serverThreadBinder.getCurrentServerSpan());
        serverThreadBinder.setCurrentSpan(null);
    }

    @Override
    public void afterCompletion(final HttpServletRequest request, final HttpServletResponse response, final Object handler, final Exception ex) {

        final ServerSpan span = (ServerSpan) request.getAttribute(HTTP_SERVER_SPAN_ATTRIBUTE);

        if (span != null) {
            serverThreadBinder.setCurrentSpan(span);
        }

        responseInterceptor.handle(new HttpServerResponseAdapter(new HttpResponse() {

            public int getHttpStatusCode() {
                return response.getStatus();
            }
        }));
    }

}
